# src/evaluator.py

from llama_index.core import VectorStoreIndex, Settings
from llama_index.llms.openai import OpenAI
from datasets import Dataset
from ragas import evaluate
from ragas.metrics import answer_relevancy, faithfulness
from typing import List

from . import config


def evaluate_query_engine(documents: List, questions: List, llm: OpenAI) -> dict:
    """
    使用给定的 LLM 和文档对查询引擎进行评估。

    Args:
        documents: 用于构建索引的文档列表。
        questions: 用于评估的问题列表。
        llm: 用于查询引擎的 LLM 实例。

    Returns:
        包含 Ragas 评估分数的字典。
    """
    print(f"正在使用模型 {llm.model} 进行评估...")

    # 设置全局参数
    Settings.llm = llm
    Settings.context_window = config.CONTEXT_WINDOW

    index = VectorStoreIndex.from_documents(documents)
    query_engine = index.as_query_engine(similarity_top_k=config.SIMILARITY_TOP_K)

    contexts = []
    answers = []

    print(f"正在处理 {len(questions)} 个评估问题...")
    for question in questions:
        response = query_engine.query(question)
        contexts.append([x.node.get_content() for x in response.source_nodes])
        answers.append(str(response))

    # 创建 Hugging Face Dataset
    ds = Dataset.from_dict({
        "question": questions,
        "answer": answers,
        "contexts": contexts,
    })

    # 使用 ragas 进行评估
    print("正在使用 Ragas 计算评估指标...")
    result = evaluate(ds, [answer_relevancy, faithfulness])
    print("评估完成。")

    return result